__doc__ =   """ Utility function and classes to manipulate entities, the functions work in-place.
            """

from numpy import std, mean, array, dot
from collections import defaultdict
from itertools import izip

from ZenExceptions import ZenException
from Einput import einput
from Misc import combinate
from Neighbors import Neighbors


def center_asu(structure, mode ='raw'):
    """ Chooses an asymmetric unit closest to the center of the primary unit cell.
        This function does not yet give the best solution, since it does not
        explore symmetry generated equivalents.
    """
    atoms = einput(structure[(0,)], 'A')                    # we work on the first model
    coords = array(atoms._data_children('coords'))
    cntr = array([mean(coords, axis =0)])                   # calculate center of mass
    fcntr = dot(cntr, structure.header[mode + '_fmx'].transpose()) % 1             # fractionalize and reduce to unit cell
    new_cntr = dot(fcntr,  structure.header[mode + '_omx'].transpose())            # orthogonalize new center of mass
    new_coords = coords + (new_cntr - cntr)                 # move coordinates to new center of mass
    for (atom_id, new_coord) in izip(atoms.keys(), new_coords):
        atoms[atom_id].coords = new_coord

def find_coordinations(entities, level ='A', xtra_key =None, cnt_key =None):
    """ Determines the chains coordinating another chain.
        @type ents: entity or list of entities
        @param ents: residue or atom entities.
        @type level: level string
        @param level: residue ('A') or atom ('R') annotation level
        @type xtra_key:
        @param xtra_key:
    """
    # The coordination number of a chain is the number of different
    # chains it is in contact with. For a single chain in the assymetric
    # unit all contacts will be generated by symmetry operations.
    # For multiple chains in the ASU each chain can have a different
    # coordination, and some of the interactions should be within the ASU
    # contacts have to be computed also within the ASU.
    entities = einput(entities, level)          # the level is arbitrary but done at the atom
    xtra_key = (xtra_key or level.lower() + 'CN')     # level elimnates the need to transfer
    cnt_key =  (cnt_key  or 'tCNT')         # contacts on residues.
    for entity in entities:
        coordination = None
        coordinations = set()
        contacts = entity.xtra.get(cnt_key)
        chain_self = entity.get_parent('C').name    # find parent chain
        if contacts: # skip get-None
            for (full_id, (dist, unit_cell, symop)) in contacts.iteritems():
                chain_id = full_id[2]
                coordination = "%s:%s:%s:%s" % (chain_self, chain_id, unit_cell, symop)
                coordinations.add(coordination)
            entity.xtra[xtra_key] = tuple(sorted(coordinations))

def clean_ical(entities, pretend =True, mask =True):
    """ Remove entities with ambiguous (i)nsertion (c)odes or (a)lternate (l)ocations.
    """
    # this function does not check for occupancy. I retains the residue which is first
    # sorting by id number, insertion code and finally name. Residues without IC come
    # first. Atoms within a retained residue are sorted according to PDB rules the first
    # one is chosen. If The first entity has an IC or alt_loc different from ' ' it
    # will be changed.
    conflicts = []
    changes = []
    residues = einput(entities, 'R')
    id_r = [[None, None, None]]
    for r in residues.sorted_values():  # sort by id, ic, name
        id_a = [[None, None]]
        if r.res_id == id_r[0][1]:      # on collision choose first ...
            conflicts.append(r.get_full_id())
            if not pretend:
                if mask:
                    r.set_masked(True)
                else:
                    r.parent.del_child(r.id)
            continue                    # an entity could be in other holders
        # keep it there as-is
        for a in r.sorted_values():     # sort by id, alt_loc (' ', 'A' ...)
            if a.at_id == id_a[0][0]:   # on collision choose first
                conflicts.append(a.get_full_id())
                if not pretend:
                    if mask:
                        a.set_masked(True)
                    else:
                        r.del_child(a.id)
            else:
                if a.id[0][1] != ' ':
                    changes.append((a.get_full_id(), ((a.id[0][0], ' '),)))
                    if not pretend:
                        a.set_alt_loc(' ')
                        try:
                            a.parent.update_ids()
                        except AttributeError:
                            pass
                id_a = a.id
        if r.id[0][2] != ' ':
            changes.append((r.get_full_id(), ((r.id[0][0], r.id[0][1], ' '),)))
            if not pretend:
                r.set_res_ic(' ')
                try:
                    r.parent.update_ids()
                except AttributeError:
                    pass
        id_r = r.id
    return (changes, conflicts)

def entities_dot_matrix(entities, matrix, level ='A', ctype ='coords'):
    atoms = einput(entities, level)
    coords = array(atoms._data_children(ctype))
    coords = dot(coords, matrix)
    for (atom, new_coords) in izip(atoms.values(),coords):
        atom.coords = new_coords

def entities_plus_vector(entities, vector, level ='A', ctype ='coords'):
    atoms = einput(entities, level)
    coords = array(atoms._data_children('coords'))
    coords = coords + vector
    for (atom, new_coords) in izip(atoms.values(),coords):
        atom.coords = new_coords

def entities_whiten(entities, level ='A', ctype ='coords' ):
    """ Whitens coordinates i.e. divides the the each coordinates
        by the std. dev. of all coordinates in each dimension).
        Useful only for clustering.
    """
    ents = einput(entities, level)
    coords = ents._data_children(ctype, forgiving =False)
    std_dev = std(coords, axis =0)
    wcoords = coords / std_dev
    for (i, ent) in enumerate(ents):
        setattr(ent, ctype, wcoords[i])



def transfer_nearest_xtra(entities, points, level, xtra_key, ctype =None, dest_key =None):
    """ Transfer xtra from nearest point(entity) for given coordinate system.
    """
    ctype = (ctype or 'coords')
    dest_key = (dest_key or xtra_key)
    ents = einput(entities, level)
    knn = Neighbors(points, level, ctype)
    for (id_, entity) in ents.iteritems():
        nn = knn.get_neighbors(entity, 1, level)
        data = nn._data_children(xtra_key, xtra =True)[0]
        entity.xtra[dest_key] = data

def transfer_contact_xtra(entities, points, level, xtra_key, cnt_key, dest_key =None):
    """ Transfer xtra from all contacts for a given contact type.
    """
    destination_key = (dest_key or xtra_key)
    atoms = einput(points, 'A')                   # contacts are A-A
    ents = einput(entities, level)
    for ent in ents:
        full_ids = ent.xtra[cnt_key].keys()       # contacting entity keys
        catoms = atoms.get_children(full_ids).values()  # contacting entities
        if catoms:
            cents = einput(catoms, level)         # transfer from right level
            ent.xtra[dest_key] = cents._data_children(xtra_key, xtra =True)



#EOF